#pragma once

#include <torch/torch.h>

class Net : public torch::nn::Module
{
private:
    int numSegement = 1;
    std::vector<float> rangeX{0, 0};
    std::vector<int> nearestIdxs;
    at::Tensor inpNearSegPoi;
    torch::Tensor segPoi;
    torch::nn::Linear fit{nullptr};

public:
    Net();
    ~Net();
    Net(const int numSeg, const std::vector<float> rangeX_);

    at::Tensor forward(at::Tensor x);

    void printSegementPoint();

    void checkSegementPoint();
};

class CustomDtaset : public torch::data::Dataset<CustomDtaset>
{
private:
    torch::Tensor inputs_; ///< 输入
    torch::Tensor labels_; ///< 标签

public:
    CustomDtaset(torch::Tensor input_data, torch::Tensor label_data);
    torch::data::Example<> get(size_t index) override;
    torch::optional<size_t> size() const override;
};
